from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableWithMessageHistory

from a0base.base_llm import pop_llm

# 定义 Prompt
prompt = ChatPromptTemplate.from_messages(("human", "{input}"))
parser = StrOutputParser()
# 构建处理链：将提示词模板、语言模型和输出解析器组合
chain = prompt | pop_llm | parser

# 创建内存聊天历史记录实例，用于存储对话历史
history = InMemoryChatMessageHistory()
# 创建带消息历史的可运行对象，用于处理带历史记录的对话
runnable = RunnableWithMessageHistory(chain, get_session_history=lambda: history)

# 清空历史记录
history.clear()
# 调用模型处理第一次输入，询问模型身份
print(runnable.invoke({"input": "你是谁？"}))
# 调用模型处理第二次输入，要求重复上一次回答
print(runnable.invoke({"input": "再重复一次"}))
